__all__ = ['FaceDataset', 'random_walk']

import random

import numpy as np
from PIL import Image


def random_walk(mu, x0, sigma_2, n):
    """
    计算随机游走向量
    """

    # 向量初始化
    result_0 = [x0] * n
    while True:
        # 生成同分布的n维正态分布随机向量
        wt = [random.normalvariate(0, sigma_2) for i in range(n)]
        # 计算随机游走向量
        result = [mu + result_0[i] + wt[i] for i in range(n)]
        yield result
        # 保存上一步结果
        result_0 = result

class FaceDataset:
    """
    脸部数据集迭代器
    """

    def __init__(self, path_list):
        self.path_list = path_list  # 图片路径列表
        self._now = 0               # 已加载指针，每加载一张图片，指针+1
        self.images = []            # 储存已加载的图片数组

    def clean(self):
        """
        清除所有已加载的图片，使已加载指针归零
        """
        self.images = []
        self._now = 0

    @staticmethod
    def image2ndarray(image: Image.Image):
        return np.asarray(image)

    def __iter__(self):
        """
        先清除所有已加载信息
        """
        self.clean()
        return self

    def __len__(self):
        """
        对象长度定义为图片路径的数量
        """
        return len(self.path_list)

    def __next__(self):
        if self._now < self.__len__():
            img = Image.open(self.path_list[self._now])
            na = self.image2ndarray(img)
            self.images.append(na)
            self._now += 1
            return na
        else:
            raise StopIteration(f"{self.__len__()}个路径已全部加载完毕")
